#
#   Perturbation1D.py
#
#   A Library to perturb the system's potential energy and wavefunction.
#

#
#   Copyright (C) 2012
#
#   This program is free software: you can redistribute it and/or modify
#   it under the terms of the GNU General Public License as published by
#   the Free Software Foundation, either version 3 of the License, or
#   (at your option) any later version.
#
#   This program is distributed in the hope that it will be useful,
#   but WITHOUT ANY WARRANTY; without even the implied warranty of
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#   GNU General Public License for more details.
#
#   You should have received a copy of the GNU General Public License
#   along with this program.  If not, see <http://www.gnu.org/licenses/>.
#


# -------------------------------------------------------------------------------

# Define / Import our Simple Harmonic Oscillator

from HarmonicOscillator1D import V as Unperturbed_V
from HarmonicOscillator1D import make_eigenfunction_n
from HarmonicOscillator1D import hbar
from HarmonicOscillator1D import omega
from cmath import exp # for make_eigenfunction

# Start by defining the known, unperturbed quantities:

Unperturbed_Psis = [] # unperturbed wavefunctions n = 0, 1, 2, 3, ...
Unperturbed_Energies = []

for n in range(150):
    Unperturbed_Psis.append(eval(make_eigenfunction_n(n)))
    Unperturbed_Energies.append(hbar * omega * (n + 0.5))


# Beyond this point, the code does not assume we are perturbing a Simple Harmonic Oscillator.
# It should be possible to perturb any system, changing only the definitions of
#   (i)   Unperturbed_V(x) ,
#   (ii)  Unperturbed_Psis[](t,x) , and
#   (iii) Unperturbed_Energies[]
# above.

# ----------------------------------------------------------------------------

# Define the peturbation (some time-independent function)

def H_prime(x):
    return 2.0 * x + 0.7 * x**3

epsilon = 0.03 # perturbation strength

# Use this to redefine the classical potential
# (This is the potential needed for SchrodingerNDT)

V = lambda x: Unperturbed_V(x) + epsilon*H_prime(x)

# Terms of the power series are stored in lists of lists.

Energies = [[]]
Wvfnctns = [[]] 

for n in range(150):
    Energies.append([])
    Wvfnctns.append([])

# We get j-th order correction of nth energy level with Energies[n][j]
# also use Wvfnctns[n][j] to get jth order correction to nth eigenstate.
# can get the approximated energy of the nth state as sum(Energies[n]), and similarly for Wvfnctns

# These data structures will be filled up after defining
# the iterative method needed to get the corrections (below)


# ----------------------------------------------------------------------------

global_time = 0.0
# Perturbation theory uses the Time-Independent Schrodinger Equation.
# The time-dependence of the perturbed eigenfunctions is added by solving
# the Time Independent Schrodinger Equation for each of them, which requires
# that we have the j-th order energy correction. We must wait to add the
#
#   exp(-i*(omega(j))*t)
#
# term at the very end, and work at a global time until then.
# (we chose global_time == 0.0 for convenience).

def segments(min,max,steps):
    return [i/float(steps)*(max-min)+min for i in range(steps+1)]

def identity (x):
    return 1

def memoize(f):
    # memoization is a functional programming trick,
    # used when the function f is
    #   (i)  expensive to compute, and
    #   (ii) referentially transparent
    #
    # Basically, we store calculated values of f in a table
    # and whenever f is ever called with the same input parameters,
    # we can return the stored value instead of recalculating it.
    class memodict(dict):
        def __init__(self, f):
            self.f = f
        def __call__(self, *args):
            return self[args]
        def __missing__(self, key):
            ret = self[key] = self.f(*key)
            return ret
    return memodict(f)

comb = segments(-40,40,5000)
step = comb[10]-comb[9]

@memoize
def overlap(bra_state,perturbation,ket_state,comb=comb):
    # Calculate the `Matrix Elements' and overlap integrals of perturbation theory.
    # This is just numerical integration.
    return sum([bra_state(global_time,i).conjugate() * ket_state(global_time,i) * perturbation(i) for i in comb]) * step

@memoize
def a(j,m,n): # the overlap < m^(0) | n^(j) >
    if m==n:
        #use normalisation of the perturbed state
        if j == 1:
            return 0.0 # analytically.
        elif j > 1:
            return -0.5 * sum([abs(a(j-k,k,n))**2 for k in range(1,j)]) # runs ups to j-1
        else:
            print "j <= 0 should never be calculated. Aborting.\n\n"
            exit(0)
    elif j==1: # use the simple formula
        return - overlap(Unperturbed_Psis[m],H_prime,Unperturbed_Psis[n]).real / (Unperturbed_Energies[m] - Unperturbed_Energies[n])
    elif j>1: # actually do the work
        denominator = Unperturbed_Energies[m] - Unperturbed_Energies[n]
        numerator = - overlap(Unperturbed_Psis[m],H_prime,Unperturbed_Psis[j-1])
        # now remove the remaining couple of terms:
        for k in range(1,j): # range runs up to j-1...
            numerator += Energies[n][k] * overlap(Unperturbed_Psis[m], identity, Wvfnctns[n][j-k])
        # finally
        return (numerator / denominator).real
    else:
        print "j <= 0 should never be calculated. Aborting.\n\n"
        exit(0)

# We now have enough formulae to express the jth order corrections to the nth states:

def make_Energy_correction(j,n):
    correction = overlap(Unperturbed_Psis[n],H_prime,Wvfnctns[n][j-1])
    for k in range(1,j): # runs up to j-1
        correction -= Energies[n][k] * overlap(Wvfnctns[n][0],identity,Wvfnctns[n][j-k])
    return correction.real * epsilon**j

def make_Psi_correction(j,n):
    highest_mixing = n+9
    # Keep this range as small as possible in order to run fast!
    m_values = range(0,highest_mixing)
    return lambda t,x: sum([a(j,m,n)*Unperturbed_Psis[m](global_time,x) for m in m_values]) * epsilon**j


# ----------------------------------------------------------------------------------------

# We now fill up the data structures Energies and Wvfnctns, iteratively

# Unperturbed system is 0th order:

for n in range(150):
    Energies[n].append(Unperturbed_Energies[n]) # nth eigenvalue
    Wvfnctns[n].append(Unperturbed_Psis[n]) # nth eigenfucntion


from math import log
from math import ceil
j_max = int(ceil(log(1e-9) / log (epsilon) - 1))
# Guarantees that the remainder is of order 10^-9
# i.e. we never have to worry about truncating the infinite sum

for n in range(3): # can go up to 150... but only if we actually need to go that high, this can be very long to calculate!
    for j in range(1,j_max+1): # do not overwrite the 0th order correction
        print("Calculating corrections to the n=="+str(n)+" eigenstate at order j == "+str(j))
        Energies[n].append(make_Energy_correction(j,n))
        Wvfnctns[n].append(make_Psi_correction(j,n))


# ----------------------------------------------------------------------------------------

# We can now build the perturbed eigenfunctions, with their time-dependences:

def make_Perturbed_Psi_n(n):
    def Perturbed_Eigenfunction (t,x,n=n):
        return sum([wvfn(global_time,x) for wvfn in Wvfnctns[n]]) * exp((0-1j)*sum(Energies[n])*t/hbar)
    return Perturbed_Eigenfunction
    # we need to pass n=n as an argument to avoid using the evaluation scope n instead of the definition scope n.
    # Actually, I'm not quite sure why we need to do this, but it works now and it didn't before so that's that.


# ----------------------------------------------------------------------------------------

# Let's create a few things to export:

PPsi1 = make_Perturbed_Psi_n(1)
PPsi2 = make_Perturbed_Psi_n(2)

from math import sqrt

def Psi(t,x):
    return (PPsi1(t,x)+PPsi2(t,x))/sqrt(2)

# ----------------------------------------------------------------------------------------


if __name__ == "__main__":

    print("\n\nDiagnostic Tests...\n")
    
    # Define ALL the corrected eigenstates:

    Perturbed_Eigenfunctions = []
    for n in range(13):
        j = lambda t,x,n=n: sum([wvfn(global_time,x) for wvfn in Wvfnctns[n]]) * exp((0-1j)*sum(Energies[n])*t/hbar)
        Perturbed_Eigenfunctions.append(make_Perturbed_Psi_n(n))
    
    print("Values: (shouldn't be too far away from 0)")
    for n in range(2):
        print(abs(Perturbed_Eigenfunctions[n](0,1.0) - Unperturbed_Psis[n](0,1.0)))
    print("Try again at a different time...")
    for n in range(2):
        print(abs(Perturbed_Eigenfunctions[n](0.23,1.0) - Unperturbed_Psis[n](0.23,1.0)))

    print("\nOrthonormality:\n")
    for n in range(2):
        for m in range(n+1):
            print(overlap(Perturbed_Eigenfunctions[m],identity,Perturbed_Eigenfunctions[n]).real)

    print("\nMixing Coefficients:\n")
    for n in range(12):
        for m in range(max(0,n-10),n+11):
            for j in range(1,j_max+1):
                print( "a"+str([j,m,n])+ " = "+ str(a(j,m,n)))